/**
* sum_matrix
*/
#include <cuda_runtime.h>
#include <cstdio>
#include <freshman.h>

void sumMatrix2D_CPU(float *A, float *B, float *C, const int nx, const int ny) {
    float *a = A;
    float *b = B;
    float *c = C;
    for (int j = 0; j < ny; j++) {
        for (int i = 0; i < nx; i++) {
            c[i] = a[i] + b[i];
        }
        a += nx;
        b += nx;
        c += nx;
    }
}

__global__ void sumMatrix2D(float *A, float *B, float *C, const int nx, const int ny) {
    int ix = threadIdx.x + blockIdx.x * blockDim.x;
    int iy = threadIdx.y + blockIdx.y * blockDim.y;
    int idx = ix + iy * nx;

    if (ix < nx && iy < ny) {
        C[idx] = A[idx] + B[idx];
    }
}

int main() {
    printf("starting.....\n");
    initDevice(0);

    int nx = 1 << 10;
    int ny = 1 << 12;
    int nxy = nx * ny;
    int nBytes = nxy * sizeof(float);

    //malloc
    float *A_host = (float *) malloc(nBytes);
    float *B_host = (float *) malloc(nBytes);
    float *C_host = (float *) malloc(nBytes);
    float *C_from_GPU = (float *) malloc(nBytes);

    //initData
    initData(A_host, nxy);
    initData(B_host, nxy);

//    printMatrix(A_host, nx, ny);

    float *A_dev = NULL;
    float *B_dev = NULL;
    float *C_dev = NULL;
    //cuda malloc
    CHECK(cudaMalloc((float **) &A_dev, nBytes));
    CHECK(cudaMalloc((float **) &B_dev, nBytes));
    CHECK(cudaMalloc((float **) &C_dev, nBytes));

    // 把数据给GPU
    CHECK(cudaMemcpy(A_dev, A_host, nBytes, cudaMemcpyHostToDevice));
    CHECK(cudaMemcpy(B_dev, B_host, nBytes, cudaMemcpyHostToDevice));

    //cou compute
//    cudaMemcpy(C_from_GPU,C_dev,nBytes,cudaMemcpyDeviceToHost); //TODO why?
    double iStart = cpuSecond();

    sumMatrix2D_CPU(A_host, B_host, C_host, nx, ny);
    double iElapse = cpuSecond() - iStart;
    printf("CPU execution time elapse %f sec\n", iElapse);

    int dimx = 32, dimy = 32;

    //2d block and 2d grid
    dim3 block_0(dimx, dimy);
    dim3 grid_0((nx - 1) / block_0.x + 1, (ny - 1) / block_0.y + 1);
    iStart = cpuSecond();
    sumMatrix2D<<<grid_0, block_0>>>(A_dev, B_dev, C_dev, nx, ny);
    CHECK(cudaDeviceSynchronize());
    iElapse = cpuSecond() - iStart;
    printf("GPU Execution configuration<<<(%d,%d),(%d,%d)>>> Time elapsed %f sec\n",
           grid_0.x, grid_0.y, block_0.x, block_0.y, iElapse);

    CHECK(cudaMemcpy(C_from_GPU, C_dev, nBytes, cudaMemcpyDeviceToHost));
    checkResult(C_host, C_from_GPU, nxy);

    // 1d block and 1d grid
    dimx = 32;
    dim3 block_1(dimx);
    dim3 grid_1((nxy - 1) / block_1.x + 1);
    iStart = cpuSecond();
    sumMatrix2D<<<grid_1, block_1>>>(A_dev, B_dev, C_dev, nx , ny);
    CHECK(cudaDeviceSynchronize());
    iElapse = cpuSecond() - iStart;
    printf("GPU Execution configuration<<<(%d,%d),(%d,%d)>>> Time elapsed %f sec\n",
           grid_1.x, grid_1.y, block_1.x, block_1.y, iElapse);
    CHECK(cudaMemcpy(C_from_GPU, C_dev, nBytes, cudaMemcpyDeviceToHost));
    checkResult(C_host, C_from_GPU, nxy);

    // 2d block and 1d grid
    dimx = 32;
    dim3 block_2(dimx);
    dim3 grid_2((nx - 1) / block_2.x + 1, ny);
    iStart = cpuSecond();
    sumMatrix2D<<<grid_2, block_2>>>(A_dev, B_dev, C_dev, nx, ny);
    CHECK(cudaDeviceSynchronize());
    iElapse = cpuSecond() - iStart;
    printf("GPU Execution configuration<<<(%d,%d),(%d,%d)>>> Time elapsed %f sec\n",
           grid_2.x, grid_2.y, block_2.x, block_2.y, iElapse);
    CHECK(cudaMemcpy(C_from_GPU, C_dev, nBytes, cudaMemcpyDeviceToHost));
    checkResult(C_host, C_from_GPU, nxy);
}